#include <argp.h>
#include <bpf/libbpf.h>
#include <sys/resource.h>
#include <kprobe_test.h>
#include <stdio.h>
#include <signal.h>
#include <kprobe_test.skel.h>
#include <cmdparser.h>

static cmdp_action_t head_process(cmdp_process_param_st* paras);

struct {
	bool kstack;
	bool ustack;
} arg_kprobe_test;

cmdp_command_st kprobe_test_cmd = {
	.name = "kprobe_test",
	.desc = "kprobe get kstack & ustack.",
	.doc = "kprobe_test [<args>]\n"
		   "\n",
	.options =
		(cmdp_option_st[]){
			{'k', "kstack", "Show kstack info.", CMDP_TYPE_BOOL, &arg_kprobe_test.kstack},
			{'u', "ustack", "Show kstack info.", CMDP_TYPE_BOOL, &arg_kprobe_test.ustack},
			{0},
		},
	.fn_process = head_process,
};

typedef struct {
	struct ring_buffer* rb;
	struct kprobe_test_bpf* skel;
} kprobe_test_obj;

static kprobe_test_obj kprobe_obj = {0};
static bool flag = false;

static cmdp_action_t head_process(cmdp_process_param_st* paras) {
	if (paras->opts) {
		flag = true;
		return CMDP_ACT_OK;
	}

	return CMDP_ACT_SHOW_HELP;
}

static int libbpf_print_fn(enum libbpf_print_level level, const char* format, va_list args) {
	return vfprintf(stderr, format, args);
}

static volatile bool exiting = false;

static void sig_handler(int sig) {
	exiting = true;
}

static int handle_event(void* ctx, void* data, size_t data_sz) {
	struct kprobe_test_event* event = data;

	if (event->kstack_sz <= 0 && event->ustack_sz <= 0)
		return 1;

	printf("COMM: %s (pid=%d) @ CPU %d\n", event->comm, event->pid, event->cpu_id);

	if (arg_kprobe_test.kstack && event->kstack_sz > 0) {
		printf("Kernel:\n");

		for (int i = 0; i < event->kstack_sz / sizeof(__u64); i++) {
			printf("0x%llx\n", event->kstack[i]);
		}
	} else {
		printf("No Kernel Stack\n");
	}

	if (arg_kprobe_test.ustack && event->ustack_sz > 0) {
		printf("Userspace:\n");
		for (int i = 0; i < event->ustack_sz / sizeof(__u64); i++) {
			printf("0x%llx\n", event->ustack[i]);
		}
	} else {
		printf("No Userspace Stack\n");
	}

	printf("\n");
	return 0;
}

int kprobe_test_release() {
	struct kprobe_test_bpf* skel = kprobe_obj.skel;
	struct ring_buffer* rb = kprobe_obj.rb;

	/* Clean up */
	if (rb)
		ring_buffer__free(rb);
	if (skel)
		kprobe_test_bpf__destroy(skel);

	return 0;
}

void kprobe_test_run() {
	struct ring_buffer* rb = kprobe_obj.rb;
	int err;

	/* Process events */
	printf("%-8s %-5s %-16s %-7s %-7s %s\n", "TIME", "EVENT", "COMM", "PID", "PPID",
		   "FILENAME/EXIT CODE");
	while (!exiting) {
		err = ring_buffer__poll(rb, 100 /* timeout, ms */);
		/* Ctrl-C will cause -EINTR */
		if (err == -EINTR) {
			err = 0;
			break;
		}
		if (err < 0) {
			printf("Error polling perf buffer: %d\n", err);
			break;
		}
	}
}

int kprobe_test_open(int argc, char* argv[]) {
	struct kprobe_test_bpf* skel;
	struct ring_buffer* rb;
	int err = 0;

	if (cmdp_run(argc, argv, &kprobe_test_cmd, NULL) || !flag)
		return -1;
	/* Set up libbpf errors and debug info callback */
	libbpf_set_print(libbpf_print_fn);

	/* Cleaner handling of Ctrl-C */
	signal(SIGINT, sig_handler);
	signal(SIGTERM, sig_handler);

	/* Load and verify BPF application */
	skel = kprobe_test_bpf__open();
	if (!skel) {
		fprintf(stderr, "Failed to open and load BPF skeleton\n");
		return -1;
	}
	kprobe_obj.skel = skel;
	/* Parameterize BPF code with minimum duration parameter */
	skel->rodata->min_duration_ns = 1000000ULL;

	/* Load & verify BPF programs */
	err = kprobe_test_bpf__load(skel);
	if (err) {
		fprintf(stderr, "Failed to load and verify BPF skeleton\n");
		goto cleanup;
	}

	/* Attach tracepoints */
	err = kprobe_test_bpf__attach(skel);
	if (err) {
		fprintf(stderr, "Failed to attach BPF skeleton\n");
		goto cleanup;
	}

	/* Set up ring buffer polling */
	rb = ring_buffer__new(bpf_map__fd(skel->maps.cb_events), handle_event, NULL, NULL);
	if (!rb) {
		err = -1;
		fprintf(stderr, "Failed to create ring buffer\n");
		goto cleanup;
	}
	kprobe_obj.rb = rb;
	return 0;

cleanup:
	kprobe_test_release();
	return -1;
}
